#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import re
import tensorflow as tf


class EmbeddingVariableOption:
    """ option for embedding service table. """

    def __init__(self, filter_option=None,
                 padding_option=None,
                 evict_option=None,
                 completion_option=None,
                 storage_option=None,
                 feature_freezing_option=None,
                 communication_option=None):
        self.filter_option = filter_option
        self.padding_option = padding_option
        self.evict_option = evict_option
        self.completion_option = completion_option
        self.storage_option = storage_option
        self.feature_freezing_option = feature_freezing_option
        self.communication_option = communication_option


class CounterFilter:
    """ Counter filter for embedding table. """

    def __init__(self, filter_freq, default_key_or_value, default_key=None, default_value=None):
        self.filter_freq = filter_freq
        self.default_key = default_key
        self.default_value = default_value
        self.default_key_or_value = default_key_or_value


class PaddingParamsOption:
    """ padding key option for embedding service table. """

    def __init__(self, padding_key=None,
                 mask=True):
        self.padding_key = padding_key
        self.mask = mask


class CompletionKeyOption:
    """ completion key option for embedding service table. """

    def __init__(self, completion_key=None, mask=1):
        self.completion_key = completion_key
        self.mask = mask


class EvictOption:
    """ Evict option for embedding table. """

    def __init__(self, steps_to_live):
        self.steps_to_live = steps_to_live


def check_common_init_params(name, init_vocabulary_size, embedding_dim, embedding_type, mask_zero):
    if (name is None) or (init_vocabulary_size is None) or (embedding_dim is None):
        raise ValueError("table name, init_vocabulary_size and embedding_dim can not be None.")
    if not isinstance(name, str):
        raise TypeError("embedding table name must be string.")
    regex = re.compile('[@!#$%^&*()<>?/\|}{~:]')
    if regex.search(name) is not None:
        raise ValueError("table name contains illegal character.")
    if (not isinstance(init_vocabulary_size, int)) or (not isinstance(embedding_dim, int)):
        raise ValueError("init_vocabulary_size and embedding_dim must be int.")
    if init_vocabulary_size < 0:
        raise ValueError("init_vocabulary_size can not be smaller than zero.")
    if embedding_dim <= 0:
        raise ValueError("embedding_dim must be greater than zero.")
    if (embedding_type != "PS") and (embedding_type != "data_parallel"):
        raise TypeError("embedding_type must be PS or data_parallel")
    if not isinstance(mask_zero, bool):
        raise TypeError("mask zero must be bool")


def check_each_initializer(initializer_mode, min_value, max_value, constant_value, mu, sigma):
    if initializer_mode == 'random_uniform':
        if (min_value is None) or (max_value is None) or \
                (not isinstance(min_value, (float, int))) or (not isinstance(max_value, (float, int))):
            raise ValueError("If initializer is random_uniform, min and max can not be None, must be int or float.")
    if initializer_mode == 'truncated_normal':
        if (min_value is None) or (max_value is None) or (mu is None) or (sigma is None) or \
                (not isinstance(min_value, (float, int))) or (not isinstance(max_value, (float, int))) or \
                (not isinstance(mu, (float, int))) or (not isinstance(sigma, (float, int))):
            raise ValueError("If initializer is truncated_normal, min, max, mu and sigma can not be None,"
                             "and they must be int or float.")
    if initializer_mode == 'constant':
        if (constant_value is None) or (not isinstance(constant_value, (float, int))):
            raise ValueError("If initializer is constant, constant_value can not be None, must be float or int.")


def check_init_params_type(key_dtype, value_dtype, init_vocabulary_size, embedding_dim, multihot_lens, allow_merge):
    if (key_dtype is not tf.int64) or (value_dtype is not tf.float32):
        raise TypeError("key_dtype only support tf.int64, value_dtype only support tf.float32 now.")
    if (not isinstance(init_vocabulary_size, int)) or (not isinstance(embedding_dim, int)) or \
            (not isinstance(multihot_lens, int)) or (not isinstance(allow_merge, bool)):
        raise TypeError("init_vocabulary_size, embedding_dim, multihot_lens must be int, allow_merge must be bool.")
